
import h5py
import matplotlib.pyplot as plt 
import numpy as np 

def process(pred):
    # b, c, d, w, h = pred.shape
    # pred_new = np.zeros([d, w, h])
    # pred_new[pred[0, 0] != 0] = 1
    # pred_new[pred[0, 1] != 0] = 2
    # pred_new[pred[0, 2] != 0] = 3
    # return pred_new 
    return pred

data = h5py.File("./logs_liver/diffusion_plots/liver_plot8.h5")

## 25-s186
## 8-s196
## 25 168
name = "8-s196"
index = 196

save_dir=f"./logs_liver/diffusion_saveplots/{name}/"
import os 
os.makedirs(save_dir, exist_ok=True)
from PIL import Image

def save_image(image_array, file_name):
    mat = np.array(image_array)
    mat = mat.astype(np.uint8)
    dst = Image.fromarray(mat, 'P')
    dst.save(file_name)

def save_index_image(data, i):
    save_image(data["image"][0, 0, i, :, :], os.path.join(save_dir, "image.png"))
    save_image(process(data["label"])[0, 0, i, :, :], os.path.join(save_dir, "label.png"))

    for key in data.keys():
        if "image" not in key and "label" not in key:

            index_data = process(data[key])[0, i, :, :]
            save_image(index_data, os.path.join(save_dir, f"{key}.png"))

    # save_image(process(data["deepunet_v2"])[:, :, i], os.path.join(save_dir, "deepunet_v2.png"))
    # save_image(process(data["segresnet"])[:, :, i], os.path.join(save_dir, "segresnet.png"))
    # save_image(process(data["unetr"])[:, :, i], os.path.join(save_dir, "unetr.png"))
    # save_image(process(data["swinunetr"])[:, :, i], os.path.join(save_dir, "swinunetr.png"))
    # save_image(process(data["transvw"])[:, :, i], os.path.join(save_dir, "transvw.png"))
    # save_image(process(data["modelgenesis"])[:, :, i], os.path.join(save_dir, "modelgenesis.png"))
    # save_image(process(data["unetformer"])[:, :, i], os.path.join(save_dir, "unetformer.png"))

if __name__ == "__main__":


    print(data)
    print(data.keys())

    save_index_image(data, index)

    # print("hello world")
    # data_image = data["image"]
    # print(data_image.shape)
    # row = 3
    # col = 4
    # for i in range(0, data_image.shape[2]):
    #     label_i = (data["label"][0, 0, i, :, :] == 2)
    #     if label_i.sum() < 20:
    #         continue

    #     print(i)

    #     plt.subplot(row, col, 1)
    #     plt.axis('off')
    #     plt.imshow(data_image[0, 0, i, :, :], cmap="gray")
    #     plt.subplot(row, col, 2)
    #     plt.axis('off')
    #     # print(data["label"].shape)
    #     plt.imshow(process(data["label"])[0, 0, i, :, :], cmap="gray")

    #     index = 3
    #     for key in data.keys():
    #         if "image" not in key and "label" not in key:
    #             plt.subplot(row, col, index)
    #             plt.title(key)
    #             plt.axis('off')

    #             index_data = process(data[key])[0, i, :, :]
    #             plt.imshow(index_data, cmap="gray")
                

    #             index += 1
    #     plt.show()


